#!/usr/bin/env python
# --coding:utf-8--
"""
@File    : Deal_Vul.py
@Time    : 2021/8/18 15:37
@Author  : cpl
"""
from neo4j import GraphDatabase
import requests
import json
import GlobalVar
import time
from Crawl import  Crawl

class Vul:
    def __init__(self):
        self.crawl = Crawl()
        self.validvul_num = 0

    def close(self):
        self.crawl.close()

    def load_singlevul(self, tx, vul_info):
        b_withDrawnMark = vul_info['advisory']['withdrawnAt']  # b开头变量表示来自于github数据库的
        #   判断该漏洞是否已过时废弃，若未废弃则存储各项信息，若废弃，则只存储项目名
        if b_withDrawnMark is None:
            self.validvul_num += 1
            #   判断ghasid是否存在，若已存在，则说明有多个受影响范围，只更新该属性；若不存在，则新增，并把各项属性存入
            pkg_name = vul_info['package']['name']
            vul_dbid = vul_info['advisory']['databaseId']
            vul_ghsaId = vul_info['advisory']['ghsaId']
            vul_summary = vul_info['advisory']['summary']
            vul_description = vul_info['advisory']['description']
            vul_severity = vul_info['severity']
            references = []
            urls = vul_info['advisory']['references']  # 是字典，需要转换为字符串
            for k in urls:
                references.append(k['url'])

            query = (
                "MERGE (p1:Vul{ id: $vul_dbid }) "
                "SET p1.pkg = $pkg_name, p1.ghsaId = $vul_ghsaId, p1.summary = $vul_summary, p1.description = $vul_description, p1.severity = $vul_severity, p1.references = $references "
                "RETURN p1"
            )
            result = tx.run(query, vul_dbid=vul_dbid, pkg_name=pkg_name, vul_ghsaId=vul_ghsaId,
                            vul_summary=vul_summary, vul_description=vul_description, vul_severity=vul_severity, references=references)
            GlobalVar.add_ordernum()
            vulnerableVersionRange = vul_info['vulnerableVersionRange'].replace(',', ' ')
            return vul_dbid, pkg_name, vulnerableVersionRange


    def linkVulPkg(self, tx, vul_dbid, pkg_name, vulnerableVersionRange):
        self.crawl.Verutils.set_name_version(pkg_name)
        self.crawl.Verutils.load_pkgjson()
        infectedVers = self.crawl.Verutils.satisfiedVerions(vulnerableVersionRange)
        query = (
            "MATCH (p:Pkg{id:$pkg_id}) "
            "MATCH (v:Vul{id:$vul_dbid}) "
            "MERGE (v)-[:Infect]->(p) "
            "RETURN p, v"
        )
        for v in infectedVers:
            pkg_id = pkg_name + '@' + v
            tx.run(query, pkg_id=pkg_id, vul_dbid=vul_dbid)
            GlobalVar.add_ordernum()

    def createVuls(self, batch_size):
        urlq = 'https://api.github.com/graphql'
        api_token = "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"
        headersq = {'Authorization': 'token %s' % api_token}
        jsonq = {
            "query": "query { securityVulnerabilities(ecosystem: NPM, first: 20) {nodes {advisory {"
                     "databaseId ghsaId summary description references{ url } withdrawnAt } package { name } severity "
                     "vulnerableVersionRange } } }"}
        r = requests.post(url=urlq, json=jsonq, headers=headersq)

        order_num = 0
        st = time.time()
        vul_num = 0

        if r.status_code == 200:  # 判断API是否成功连接
            session = self.crawl.app.driver.session()
            tx = session.begin_transaction()

            rs = r.text  # 将返回的response类型转换为字符串
            rd = json.loads(rs)

            record = open("vulRecord.txt", "a")
            for j in rd['data']['securityVulnerabilities']['nodes']:
                curtime = time.time()
                vul_num += 1
                vul_dbid, pkg_name, vulnerableVersionRange = self.load_singlevul(tx, j)
                self.crawl.crawl_vul_pkg_batch(tx, pkg_name, batch_size)
                self.linkVulPkg(tx, vul_dbid, pkg_name, vulnerableVersionRange)
                if GlobalVar.get_ordernum() >= batch_size:
                    tx.commit()
                    t = time.time() - curtime
                    curtime = time.time()
                    order_num += GlobalVar.get_ordernum()
                    print("-------commit at vul-%d <%s>, %d orders cost %f s-------" % (vul_num, j['advisory']['summary'], GlobalVar.get_ordernum(), t))

                    GlobalVar.clear_ordernum()
                    session.close()
                    session = self.crawl.app.driver.session()
                    tx = session.begin_transaction()
                print("---------------finish store vul-%d <%s> ---------------" % (vul_num, j['advisory']['summary']))
                record.write("---------------finish store vul-%d <%s> ---------------\n" % (vul_num, j['advisory']['summary']))


            tx.commit()
            total_time = time.time() - st
            print("FINISH %d/%d VULS, %d INSERT ORDERS,  CONSUMING %fs"%(self.validvul_num, vul_num, order_num, total_time))
            record.write("FINISH %d VULS, %d INSERT ORDERS,  CONSUMING %fs"%(vul_num, order_num, total_time))
            record.close()
            self.close()
        else:
            print("api error!")